import os

from Crypto.PublicKey import RSA

def generate_rsa_keypair(key_size=2048):
    """
    生成 RSA 密钥对。
    :param key_size: 密钥长度，推荐 2048 或 4096。
    :return: (private_key_pem, public_key_pem)，均为字节串。
    """
    key = RSA.generate(key_size)
    private_pem = key.export_key(format='PEM')
    public_pem  = key.publickey().export_key(format='PEM')
    return public_pem, private_pem

def work():
    base_dir = os.path.dirname(os.path.abspath(__file__)).replace('\\', '/') + "/../../"
    save_path = base_dir + "data/client/rsa/"
    required_files = ['pk1.pem', 'sk1.pem', 'pk2.pem', 'sk2.pem']
    missing_files = [f for f in required_files if not os.path.isfile(os.path.join(save_path, f))]
    if  missing_files:
        pk1,sk1 = generate_rsa_keypair()
        pk2,sk2 = generate_rsa_keypair()
        with open(save_path + "pk1.pem", "wb") as f: f.write(pk1)
        with open(save_path + "sk1.pem", "wb") as f: f.write(sk1)
        with open(save_path + "pk2.pem", "wb") as f: f.write(pk2)
        with open(save_path + "sk2.pem", "wb") as f: f.write(sk2)
    print("RSA key pair generated successfully.")
    return generate_rsa_keypair()
if   __name__ == '__main__':
    work()